import os
import asyncio
import websockets
import uuid
import json
import gzip
import copy
import tempfile
from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
from dotenv import load_dotenv
import traceback

# 加载环境变量
load_dotenv()
VOLCANO_APPID = os.getenv("VOLCANO-APPID")
VOLCANO_TOKEN = os.getenv("VOLCANO-TOKEN")

app = Flask(__name__)
CORS(app)  # 启用CORS支持

# WebSocket API 配置
host = "openspeech.bytedance.com"
api_url = f"wss://{host}/api/v1/tts/ws_binary"
cluster = "volcano_tts"  # 根据官方示例，这是固定值

# version: b0001 (4 bits)
# header size: b0001 (4 bits)
# message type: b0001 (Full client request) (4bits)
# message type specific flags: b0000 (none) (4bits)
# message serialization method: b0001 (JSON) (4 bits)
# message compression: b0001 (gzip) (4bits)
# reserved data: 0x00 (1 byte)
default_header = bytearray(b'\x11\x10\x11\x00')

# 用于解析响应的常量
MESSAGE_TYPES = {11: "audio-only server response", 12: "frontend server response", 15: "error message from server"}
MESSAGE_TYPE_SPECIFIC_FLAGS = {0: "no sequence number", 1: "sequence number > 0",
                              2: "last message from server (seq < 0)", 3: "sequence number < 0"}
MESSAGE_SERIALIZATION_METHODS = {0: "no serialization", 1: "JSON", 15: "custom type"}
MESSAGE_COMPRESSIONS = {0: "no compression", 1: "gzip", 15: "custom compression method"}

async def generate_tts(text, voice_type, speed_ratio=1.0, encoding="mp3"):
    """调用火山引擎WebSocket API生成语音"""
    
    # 检查API密钥是否已设置
    if not VOLCANO_APPID or not VOLCANO_TOKEN:
        raise ValueError("火山引擎API密钥未配置，请在.env文件中设置VOLCANO-APPID和VOLCANO-TOKEN")
    
    # 准备请求参数 - 严格按照示例格式
    request_json = {
        "app": {
            "appid": VOLCANO_APPID,
            "token": VOLCANO_TOKEN,  # 使用.env中设置的token
            "cluster": cluster
        },
        "user": {
            "uid": str(uuid.uuid4())  # 生成唯一用户ID
        },
        "audio": {
            "voice_type": voice_type,
            "encoding": encoding,
            "speed_ratio": float(speed_ratio),
            "volume_ratio": 1.0,  # 默认音量
            "pitch_ratio": 1.0,   # 默认音调
        },
        "request": {
            "reqid": str(uuid.uuid4()),  # 生成唯一请求ID
            "text": text,
            "text_type": "plain",
            "operation": "submit"  # 使用流式接口
        }
    }
    
    print(f"请求参数: {json.dumps(request_json, indent=2, ensure_ascii=False)}")
    
    # 创建临时文件用于保存音频
    with tempfile.NamedTemporaryFile(suffix=f".{encoding}", delete=False) as audio_file:
        temp_filename = audio_file.name
        
        # 准备WebSocket请求数据
        payload_bytes = str.encode(json.dumps(request_json))
        payload_bytes = gzip.compress(payload_bytes)
        full_client_request = bytearray(default_header)
        full_client_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
        full_client_request.extend(payload_bytes)
        
        print(f"请求火山引擎TTS API，文本长度: {len(text)}，音色: {voice_type}")
        
        try:
            # 发起WebSocket请求 - 注意headers格式，与示例保持一致
            header = {"Authorization": f"Bearer; {VOLCANO_TOKEN}"}
            print(f"使用认证头: {header}")
            
            async with websockets.connect(api_url, extra_headers=header, ping_interval=None) as ws:
                await ws.send(full_client_request)
                
                # 接收响应并写入临时文件
                while True:
                    try:
                        response = await ws.recv()
                        done = await parse_response(response, audio_file)
                        if done:
                            break
                    except Exception as e:
                        print(f"接收响应时出错: {str(e)}")
                        raise
        except Exception as e:
            print(f"WebSocket连接失败: {str(e)}")
            traceback.print_exc()
            raise ValueError(f"火山引擎API调用失败: {str(e)}")
                    
    return temp_filename

async def parse_response(res, file):
    """解析WebSocket响应并写入音频文件"""
    # 解析协议头部
    protocol_version = res[0] >> 4
    header_size = res[0] & 0x0f
    message_type = res[1] >> 4
    message_type_specific_flags = res[1] & 0x0f
    serialization_method = res[2] >> 4
    message_compression = res[2] & 0x0f
    
    # 打印详细的响应信息用于调试
    print(f"响应: 协议版本={protocol_version:#x}, 消息类型={message_type:#x} ({MESSAGE_TYPES.get(message_type, '未知')})")
    
    payload = res[header_size*4:]
    
    # 处理音频响应
    if message_type == 0xb:  # audio-only server response
        if message_type_specific_flags == 0:  # no sequence number as ACK
            print("收到确认消息")
            return False
        else:
            sequence_number = int.from_bytes(payload[:4], "big", signed=True)
            payload_size = int.from_bytes(payload[4:8], "big", signed=False)
            audio_data = payload[8:]
            
            print(f"收到音频数据: 序列号={sequence_number}, 大小={payload_size}字节")
            
            # 写入音频数据
            file.write(audio_data)
            
            # 检查是否是最后一个消息
            if sequence_number < 0:
                print("这是最后一个数据包")
                return True
            else:
                return False
    elif message_type == 0xf:  # 错误消息
        code = int.from_bytes(payload[:4], "big", signed=False)
        msg_size = int.from_bytes(payload[4:8], "big", signed=False)
        error_msg = payload[8:]
        if message_compression == 1:
            error_msg = gzip.decompress(error_msg)
        error_msg = str(error_msg, "utf-8")
        error_detail = f"错误代码: {code}, 消息: {error_msg}"
        print(f"火山引擎TTS API错误: {error_detail}")
        raise ValueError(error_msg)
    elif message_type == 0xc:  # 前端消息
        msg_size = int.from_bytes(payload[:4], "big", signed=False)
        frontend_msg = payload[4:]
        if message_compression == 1:
            frontend_msg = gzip.decompress(frontend_msg)
        frontend_msg = str(frontend_msg, "utf-8")
        print(f"前端消息: {frontend_msg}")
        return False
    else:
        print(f"未定义的消息类型: {message_type:#x}")
        return False

def run_async(coroutine):
    """运行异步函数"""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(coroutine)
    finally:
        loop.close()

@app.route('/api/volcano-tts', methods=['POST'])
def volcano_tts():
    """处理前端请求，调用火山引擎TTS API"""
    try:
        data = request.json
        if not data:
            return jsonify({"error": "请求体为空"}), 400
            
        text = data.get("input")
        voice_type = data.get("voice", "zh_female_qingxin")
        format = data.get("response_format", "mp3")
        speed = float(data.get("speed", 1.0))
        
        if not text:
            return jsonify({"error": "文本内容为空"}), 400
        
        # 检查API密钥
        if not VOLCANO_APPID or not VOLCANO_TOKEN:
            return jsonify({"error": "火山引擎API密钥未配置，请在.env文件中设置VOLCANO-APPID和VOLCANO-TOKEN"}), 400
        
        print(f"处理火山引擎TTS请求: 音色={voice_type}, 格式={format}, 语速={speed}")
        print(f"APPID: {VOLCANO_APPID[:4]}***{VOLCANO_APPID[-4:] if len(VOLCANO_APPID) > 8 else ''}")
        print(f"TOKEN: {VOLCANO_TOKEN[:4]}***{VOLCANO_TOKEN[-4:] if len(VOLCANO_TOKEN) > 8 else ''}")
        
        # 调用WebSocket API生成语音
        try:
            audio_file = run_async(generate_tts(text, voice_type, speed, format))
            
            # 发送生成的音频文件
            return send_file(
                audio_file,
                mimetype=f'audio/{format}',
                as_attachment=True,
                download_name=f'speech.{format}'
            )
        except ValueError as ve:
            error_msg = str(ve)
            print(f"API调用出错: {error_msg}")
            return jsonify({"error": error_msg}), 400
        except Exception as e:
            import traceback
            error_trace = traceback.format_exc()
            print(f"生成语音失败: {str(e)}\n{error_trace}")
            return jsonify({"error": f"生成语音失败: {str(e)}"}), 500
    
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print(f"火山引擎TTS API错误: {str(e)}\n{error_trace}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/volcano-tts/test', methods=['GET'])
def test_volcano():
    """测试火山引擎TTS API是否正常工作"""
    try:
        if not VOLCANO_APPID or not VOLCANO_TOKEN:
            return jsonify({
                "status": "火山引擎TTS API服务正在运行",
                "api_configured": False,
                "message": "API密钥未配置，请在.env文件中设置VOLCANO-APPID和VOLCANO-TOKEN"
            })
        
        # 显示APPID和TOKEN的掩码版本，用于调试但不完全显示敏感信息
        masked_appid = VOLCANO_APPID[:4] + "****" + (VOLCANO_APPID[-4:] if len(VOLCANO_APPID) > 8 else "")
        masked_token = VOLCANO_TOKEN[:4] + "****" + (VOLCANO_TOKEN[-4:] if len(VOLCANO_TOKEN) > 8 else "")
        
        return jsonify({
            "status": "火山引擎TTS API服务正在运行",
            "api_configured": True, 
            "appid": masked_appid,
            "token": masked_token
        })
    except Exception as e:
        return jsonify({
            "status": "火山引擎TTS API服务正在运行",
            "api_configured": False,
            "error": str(e)
        })

@app.route('/api/volcano-tts/voices', methods=['GET'])
def get_voices():
    """返回火山引擎支持的音色列表"""
    voices = {
        "中文音色": [
            "zh_female_qingxin", "zh_male_qingxin", 
            "zh_female_wenzhong", "zh_male_wenzhong",
            "zh_female_huanxin", "zh_male_huanxin"
        ],
        "方言音色": [
            "zh_female_sichuan", "zh_female_dongbei"
        ],
        "英文音色": [
            "en_female_qingxin", "en_male_qingxin"
        ],
        "日文音色": [
            "ja_female_qingxin", "ja_male_qingxin"
        ]
    }
    return jsonify(voices)

if __name__ == "__main__":
    print(f"火山引擎TTS服务启动于端口5003...")
    print(f"API密钥状态: {'已设置' if VOLCANO_APPID and VOLCANO_TOKEN else '未设置'}")
    if VOLCANO_APPID and VOLCANO_TOKEN:
        print(f"APPID: {VOLCANO_APPID[:4]}***{VOLCANO_APPID[-4:] if len(VOLCANO_APPID) > 8 else ''}")
        print(f"TOKEN: {VOLCANO_TOKEN[:4]}***{VOLCANO_TOKEN[-4:] if len(VOLCANO_TOKEN) > 8 else ''}")
    print(f"测试服务: curl http://localhost:5003/api/volcano-tts/test")
    print(f"音色列表: curl http://localhost:5003/api/volcano-tts/voices")
    app.run(host='0.0.0.0', port=5003, debug=True)
